import asyncio
from abc import ABC
from functools import wraps
from typing import Type, List, Optional

from inspect import getmembers, ismethod

from langchain_core.messages import HumanMessage
from langgraph.types import interrupt
from pydantic import BaseModel


def tool_node(args_schema: Optional[Type[BaseModel]] = None, args_parser: Optional[dict[str, callable]] = None,
              stream: bool = True):
    def decorator(func):
        @wraps(func)
        async def async_wrapper(*args, **kwargs):
            return await func(*args, **kwargs)

        @wraps(func)
        def sync_wrapper(*args, **kwargs):
            return func(*args, **kwargs)

        wrapper = async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper

        wrapper.is_tool = True
        wrapper.args_schema = args_schema
        wrapper.args_parser = args_parser or {}
        wrapper.stream = stream
        return wrapper

    return decorator


class BaseNodes(ABC):

    def __init__(self, llm_chat=None, verbose=False, **kwargs):
        self.llm_chat = llm_chat
        self.verbose = verbose
        for k, v in kwargs.items():
            setattr(self, k, v)

    def get_nodes(self) -> List[dict]:
        nodes = []
        for name, method in getmembers(self, predicate=ismethod):
            if hasattr(method, 'is_tool'):
                nodes.append({
                    "name": name,
                    "func": method,
                    "args_schema": method.args_schema,
                    "args_parser": method.args_parser,
                    "stream": method.stream
                })
        return nodes

    def get_tool_methods(self) -> List[callable]:
        tool_methods = []
        for name, method in getmembers(self, predicate=ismethod):
            if hasattr(method, 'is_tool'):
                tool_methods.append(method)
        return tool_methods

    def wait_human_feedback(self, interrupt_msg: str = "interrupt", **kwargs):
        value = interrupt(interrupt_msg)
        return {
            "messages": [HumanMessage(content=value)],
            **kwargs
        }

    def node_success(self, message: str = None, **kwargs):
        if not message:
            return {**kwargs}
        if "force_goto" not in kwargs:
            kwargs["force_goto"] = None
        return {
            "messages": [HumanMessage(content=message)],
            **kwargs
        }

    def node_failure(self, message: str, **kwargs):
        if "force_goto" not in kwargs:
            kwargs["force_goto"] = None
        return {
            "messages": [HumanMessage(content=f"Result error: {message}")],
            **kwargs
        }
